module gauss

contains

    pure elemental subroutine swap(a, b)
        real(kind=8), intent(inout) :: a
        real(kind=8), intent(inout) :: b
        real(kind=8) :: tmp
        tmp = a; a = b; b = tmp
    end subroutine swap

    !> code:blocks群友提供。
    pure subroutine solve(a, b, n, m)
        ! n*n * n*m = n*n
        integer,      intent(in   ) :: n
        integer,      intent(in   ) :: m
        real(kind=8), intent(inout) :: a(n, n)
        real(kind=8), intent(inout) :: b(n, m)

        real(kind=8) :: temp, temp2
        integer      :: i, j, k, major

        ! column major method
        do i = 1, n - 1
            ! 寻找第i列主元 find major position
            major = maxloc(abs(a(i:n, i)), dim=1) + i - 1
            !数据交换,cahce miss
            call swap(a(i, i:n), a(major, i:n))
            call swap(b(i,  : ), b(major,  : ))
            if (abs(a(i, i)) < epsilon(temp)) return
            temp = 1.d0/a(i, i)
            !计算
            do j = i + 1, n
                a(i, j)       = a(i, j)*temp!归一化
                a(i + 1:n, j) = a(i + 1:n, j) - a(i + 1:n, i)*a(i, j)
            end do
            do j = 1, m
                b(i, j)       = b(i, j)*temp!归一化
                b(i + 1:n, j) = b(i + 1:n, j) - a(i + 1:n, i)*b(i, j)
            end do
            !第i列
            a(i, i)       = 1.d0
            a(i + 1:n, i) = 0.d0
        end do

        temp    = 1.d0/a(n, n)
        a(n, n) = 1.d0
        ! 回代
        do j = 1, m
            !归一化
            b(n, j) = b(n, j)*temp
            do k = 1, n - 1
                temp2         = b(n - k + 1, j)
                b(1:n - k, j) = b(1:n - k  , j) - temp2*a(1:n - k, n - k + 1)
            end do
        end do

    end subroutine solve

    !> 列主元高斯消去法，用于求解线性方程组，请自行避免矩阵奇异
    pure function linear_solve(A, b) result(ans)

        implicit none
        real(kind=8), intent(in) :: A(:, :), b(:)
        real(kind=8) :: ans(size(b))

        real(kind=8) :: M(size(A, dim=1) + 1, size(A, dim=2))
        logical      :: mask(size(A, dim=2))
        integer      :: row, column, counter, i, n, order(size(A, dim=2))

        row    = size(A, dim=1)
        column = size(A, dim=2)
        mask   = .true.
        order  = 0

        !> 构造增广矩阵
        M(:row , :) = transpose(A)
        M(row+1, :) = b

        do counter = 1, row
            i              = maxloc(abs(M(counter, :)), dim=1, mask=mask)
            mask(i)        = .false.
            order(counter) = i
            if (abs(M(counter, i)) <= epsilon(b(1))) cycle
            M(:, i) = M(:, i)/M(counter, i)
            forall (n=1:column, mask(n)) M(:, n) = M(:, n) - M(counter, n)*M(:, i)
        end do

        mask = .not. mask
        do counter = count(mask), 1, -1
            mask(order(counter)) = .false.
            forall (n=1:column, mask(n)) M(:, n) = M(:, n) - M(counter, n)*M(:, order(counter))
        end do

        ans = matmul(M(:row, :), M(row+1, :))

    end function linear_solve

    !> 列主元高斯消去法，用于求解线性方程组，请自行避免矩阵奇异
    pure function linear_solve1(A, b) result(x)

        implicit none
        real(kind=8), intent(in) :: A(:, :), b(:)
        real(kind=8) :: x(size(b))

        real(kind=8) :: M(size(A, 1), size(A, 2)+1)
        logical      :: mask(size(A, 1))
        integer      :: row, column, counter, i, n, order(size(A, 1))

        row    = size(A, 1)
        column = size(A, 2)
        mask   = .true.
        order  = 0

        !> 增广矩阵
        M(:, :column ) = A
        M(:, column+1) = b

        do counter = 1, column
            i              = maxloc(abs(M(:, counter)), dim=1, mask=mask)
            mask(i)        = .false.
            order(counter) = i
            if (abs(M(i, counter)) <= epsilon(b(1))) cycle
            M(i, :) = M(i, :)/M(i, counter)
            forall (n=1:row, mask(n)) M(n, :) = M(n, :) - M(n, counter)*M(i, :)
        end do

        mask = .not.mask
        do counter = count(mask), 1, -1
            mask(order(counter)) = .false.
            forall (n=1:row, mask(n)) M(n, :) = M(n, :) - M(n, counter)*M(order(counter), :)
        end do

        x = matmul(M(:, column+1), M(:, :column))

    end function linear_solve1

end module gauss
